package com.h.h526.controller;

import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.h.h526.pojo.PromptEntity;
import com.h.h526.pojo.R;
import jakarta.annotation.PostConstruct;
import lombok.SneakyThrows;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.*;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

@RestController
@RequestMapping
public class one {

    private final ChatClient chatClient;

    public one(OpenAiChatModel chatModel ) {
        this.chatClient = ChatClient.builder(chatModel).defaultSystem(PromptEntity.text).build();
    }

    @Autowired
    private OpenAiChatModel chatModel;

    @Autowired
    private RedisVectorStore vectorStore;
    @Autowired
    private ChatMemory chatMemory;

    @Autowired
    ChatMemoryRepository chatMemoryRepository;

    @Autowired
    private OpenAiEmbeddingModel embeddingModel;

    @PostMapping(value = "emb", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public Map emb(@RequestBody MultipartFile file) throws IOException {
        //把file转换为document
        TikaDocumentReader tdr = new TikaDocumentReader(file.getResource());
        List<Document> documents = tdr.read();
        vectorStore.add(documents);
        return Map.of("code",200);
    }
    @GetMapping("t1")
    public void t1(@RequestParam String message){
        float[] embed = embeddingModel.embed(message);
        System.out.println(embed.length);
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder().topK(1).build());

        System.out.println(documents);
    }
    @GetMapping("t3")
    public void t3(@RequestParam String message){
        float[] embed = embeddingModel.embed("傻大个是谁");
        float[] embed1 = embeddingModel.embed("傻大个是廖泉");
        System.out.println(embeddingModel.call(new EmbeddingRequest(List.of(message), EmbeddingOptionsBuilder.builder().build())));
    }
    @GetMapping("add")
    public void add(@RequestParam String message){
        vectorStore.add(List.of(Document.builder().text(message).build()));
    }
    @GetMapping(value = "t2",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> t2(@RequestParam String message){
        Document document = vectorStore.similaritySearch(message).get(0);
        QuestionAnswerAdvisor advisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(SearchRequest.builder().query(message).topK(3).build())
                .build();
        return chatClient.prompt()
                .advisors(advisor)
                .user(message)
                .stream()
                .content();
    }



    @GetMapping( value = "rag",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> t2(@RequestParam String message,@RequestParam String token) {
        System.out.println(chatMemory.get(token));
        Flux<String> content = chatClient.prompt()
                .advisors(MessageChatMemoryAdvisor.builder(chatMemory).conversationId(token).build())
                .advisors(QuestionAnswerAdvisor.builder(vectorStore)
                        .searchRequest(SearchRequest.builder().query(message).build()).build())
                .user(message)
                .stream()
                .content();
        // 使用中间变量保存最终回答内容
        AtomicReference<String> aiResponse = new AtomicReference<>("");
        chatMemory.add(token,UserMessage.builder().text(message).build());
        // 在每次收到一块响应时拼接起来
        content = getStringFlux(token, content, aiResponse);
        return content;
    }
    private @NotNull Flux<String> getStringFlux(String token, Flux<String> content, AtomicReference<String> aiResponse) {
        content = content
                .doOnNext(x -> aiResponse.accumulateAndGet(x,String::concat)) // 或 accumulateAndGet((s, t) -> s + t)
                .doOnComplete(() -> {
                    // 流结束时将完整回答加入 ChatMemory
                    String fullResponse = aiResponse.get();
                    if (fullResponse != null && !fullResponse.isEmpty()) {
                        chatMemory.add(token, new AssistantMessage(fullResponse));
                    }
                });
        return content;
    }


//    @GetMapping("/ai/add")
//    public void addTestDoc() {
//        Document doc = Document.builder()
//                .id(UUID.randomUUID().toString())
//                .text("傻大个是廖泉")
//                .metadata(Collections.singletonMap("source", "joke"))
//                .build();
//        float[] embed = embeddingModel.embed("some text!");
//        System.out.println("向量有："+embed.length+"维");
//        System.out.println(doc);
//        vectorStore.add(List.of(doc));
//    }
//
//    @GetMapping("/ai/emb")
//    public List<Document> embed1(@RequestParam(value = "message", defaultValue = "chicken") String message) {
//        List<Document> result = vectorStore.similaritySearch(
//                SearchRequest.builder()
//                        .query("傻大个是谁")
//                        .topK(5)
//                        .build()
//        );
//        System.out.println(result);
//        return result;
//    }


}
